package ssh

import (
	"fmt"
	"net"

	"golang.org/x/crypto/ssh"
)

type HandlerConfig struct {
	NewDialTcpHandler   func(srvConn *ssh.ServerConn) *DialTcpHandler
	NewListenTcpHandler func(srvConn *ssh.ServerConn) *ListenTcpHandler
	NewSessionHandler   func(srvConn *ssh.ServerConn) *SessionHandler
}

type ServerConfig struct {
	ssh.ServerConfig
	HandlerConfig
}

func NewServerConn(conn net.Conn, config *ServerConfig) (*ssh.ServerConn, error) {
	srvConn, newChs, reqs, err := ssh.NewServerConn(conn, &config.ServerConfig)
	if err != nil {
		return nil, err
	}
	dialTcpHandler := config.NewDialTcpHandler(srvConn)
	listenTcpHandler := config.NewListenTcpHandler(srvConn)
	go func() {
		srvConn.Wait()
		listenTcpHandler.Finalize()
	}()
	sessionHandler := config.NewSessionHandler(srvConn)
	go func() {
		for req := range reqs {
			switch req.Type {
			case "tcpip-forward":
				listenTcpHandler.Handle(req)
			case "cancel-tcpip-forward":
				listenTcpHandler.HandleClose(req)
			default:
				req.Reply(false, nil)
			}
		}
	}()
	go func() {
		for newCh := range newChs {
			switch newCh.ChannelType() {
			case "direct-tcpip":
				dialTcpHandler.Handle(newCh)
			case "session":
				sessionHandler.Handle(newCh)
			default:
				newCh.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %v", newCh.ChannelType()))
			}
		}
	}()
	return srvConn, nil
}
